package main

import (
	"context"
	"encoding/csv"
	"fmt"
	"github.com/milvus-io/milvus-sdk-go/v2/client"
	"github.com/milvus-io/milvus-sdk-go/v2/entity"
	"log"
	"os"
	"strconv"
	"strings"
)

func getObj(results []entity.Column, name string) entity.Column {
	var obj entity.Column
	for _, result := range results {
		if result.Name() == name {
			obj = result
		}
	}
	return obj
}

// 辅助函数：将 float32 数组转换为逗号分隔的字符串
func float32ArrayToString(array []float32) string {
	var strArray []string
	for _, val := range array {
		strArray = append(strArray, strconv.FormatFloat(float64(val), 'f', -1, 32))
	}
	return strings.Join(strArray, ",")
}

func main() {
	// 创建 Milvus 客户端
	ctx := context.Background()
	milvusClient, err := client.NewGrpcClient(ctx, "192.168.30.80:31530") // 修改为你的 Milvus 地址
	if err != nil {
		log.Fatalf("failed to create milvus client: %v", err)
	}
	defer milvusClient.Close()
	// 查询条件和需要输出的字段
	collectionName := "galleryNew"
	expr := "id > 0"
	outputFields := []string{
		"embedding", "image_type", "gallery_image_id", "path",
		"ai_model", "model_version", "prompt", "width", "height",
		"aspectRatio", "created_at", "deleted_at",
	}

	var totalRecords = 130000                                 // 根据实际记录数进行调整
	var pageSize = 1000                                       // 每页的记录数
	var totalPages = (totalRecords + pageSize - 1) / pageSize // 总页数

	// 分页获取并写入数据
	for pageNum := 0; pageNum < totalPages; pageNum++ {
		// 创建 CSV 文件，文件名递增
		fileName := fmt.Sprintf("milvus_query_results_page_%d.csv", pageNum+1)
		file, err := os.Create(fileName)
		if err != nil {
			log.Fatalf("failed to create CSV file: %v", err)
		}
		defer file.Close()

		// 初始化 CSV 写入器
		writer := csv.NewWriter(file)
		defer writer.Flush()

		// 写入表头
		writer.Write(outputFields)

		// 计算当前页的 offset
		offset := pageNum * pageSize

		// 执行查询
		opt := client.SearchQueryOptionFunc(func(option *client.SearchQueryOption) {
			option.Limit = int64(pageSize)
			option.Offset = int64(offset)
		})

		// 查询数据
		queryResults, err := milvusClient.Query(ctx, collectionName, nil, expr, outputFields, opt)
		if err != nil {
			log.Fatalf("failed to query milvus: %v", err)
		}

		// 遍历查询结果并写入 CSV
		var embeddingObj = getObj(queryResults, "embedding")
		var imageTypeObj = getObj(queryResults, "image_type")
		var galleryImageIdObj = getObj(queryResults, "gallery_image_id")
		var pathObj = getObj(queryResults, "path")
		var aiModelObj = getObj(queryResults, "ai_model")
		var modelVersionObj = getObj(queryResults, "model_version")
		var promptObj = getObj(queryResults, "prompt")
		var widthObj = getObj(queryResults, "width")
		var heightObj = getObj(queryResults, "height")
		var aspectRatioObj = getObj(queryResults, "aspectRatio")
		var createdAtObj = getObj(queryResults, "created_at")
		var deletedAtObj = getObj(queryResults, "deleted_at")

		for i := 0; i < pageSize; i++ {
			// 确保不要超过实际查询结果的数量
			if i >= int(len(queryResults)) {
				break
			}

			// 假设 embedding 是一个 []float32 数组
			embedding, _ := embeddingObj.Get(i)
			embeddingArray, ok := embedding.([]float32)
			if !ok {
				log.Fatalf("unexpected type, expected []float32")
			}
			// 将 float32 切片转换为字符串
			embeddingResult := "[" + float32ArrayToString(embeddingArray) + "]"

			imageType, _ := imageTypeObj.GetAsInt64(i)
			galleryImageId, _ := galleryImageIdObj.GetAsInt64(i)
			path, _ := pathObj.GetAsString(i)
			aiModel, _ := aiModelObj.GetAsString(i)
			modelVersion, _ := modelVersionObj.GetAsString(i)
			prompt, _ := promptObj.GetAsString(i)
			width, _ := widthObj.GetAsInt64(i)
			height, _ := heightObj.GetAsInt64(i)
			aspectRatio, _ := aspectRatioObj.GetAsDouble(i)
			createdAt, _ := createdAtObj.GetAsInt64(i)
			deletedAt, _ := deletedAtObj.GetAsInt64(i)

			record := []string{
				fmt.Sprintf("%v", embeddingResult),
				fmt.Sprintf("%v", imageType),
				fmt.Sprintf("%v", galleryImageId),
				fmt.Sprintf("%v", path),
				fmt.Sprintf("%v", aiModel),
				fmt.Sprintf("%v", modelVersion),
				fmt.Sprintf("%v", prompt),
				fmt.Sprintf("%v", width),
				fmt.Sprintf("%v", height),
				fmt.Sprintf("%v", aspectRatio),
				fmt.Sprintf("%v", createdAt),
				fmt.Sprintf("%v", deletedAt),
			}

			// 写入 CSV
			writer.Write(record)
		}

		// 刷新缓冲区，确保批量写入后立即写入文件
		writer.Flush()
		fmt.Printf("CSV 文件 %s 导出成功！\n", fileName)
	}
}
